#include "StableINSSolver3D.hh"
#include "../algorithms/Bresenham.hh"


struct AdvectionTracker {

  const bool *boundary;

  int gridSizeX, gridSizeY, gridSizeZ;

  int iEnd, jEnd, kEnd;

  AdvectionTracker(const bool *boundary, int gridSizeX, int gridSizeY, int gridSizeZ) {
    this->boundary = boundary;
    this->gridSizeX = gridSizeX;
    this->gridSizeY = gridSizeY;
    this->gridSizeZ = gridSizeZ;
    this->iEnd = -1;
    this->jEnd = -1;
    this->kEnd = -1;
  }

  inline bool operator()(int i, int j, int k) {
    if (boundary[gridSizeX * (k * gridSizeY + j) + i]) {
      return false;
    }

    iEnd = i;
    jEnd = j;
    kEnd = k;
    return true;
  }
};


StableINSSolver3D::StableINSSolver3D()
: velocityX(NULL), velocityY(NULL), velocityZ(NULL), boundary(NULL),
  vorticityX(NULL), vorticityY(NULL), vorticityZ(NULL),
  tmpGrid(NULL) {
}

StableINSSolver3D::~StableINSSolver3D() {
  freeData();
}

const float *StableINSSolver3D::getVelocityX() {
  return velocityX;
}

const float *StableINSSolver3D::getVelocityY() {
  return velocityY;
}

const float *StableINSSolver3D::getVelocityZ() {
  return velocityZ;
}

void StableINSSolver3D::resetSystem(
    int domainSizeX, int domainSizeY, int domainSizeZ, float gridSpacing,
    const float *velocityX, const float *velocityY, const float *velocityZ,
    const bool *boundary) {
  freeData();

  this->domainSizeX = domainSizeX;
  this->domainSizeY = domainSizeY;
  this->domainSizeZ = domainSizeZ;
  this->hx = gridSpacing;
  this->hy = gridSpacing;
  this->hz = gridSpacing;
  this->gridSizeX = 1 + domainSizeX + 1;
  this->gridSizeY = 1 + domainSizeY + 1;
  this->gridSizeZ = 1 + domainSizeZ + 1;
  this->gridSize = gridSizeX * gridSizeY * gridSizeZ;

  this->velocityX  = new float[gridSize];
  this->velocityY  = new float[gridSize];
  this->velocityZ  = new float[gridSize];
  this->tmpVelocityX  = new float[gridSize];
  this->tmpVelocityY  = new float[gridSize];
  this->tmpVelocityZ  = new float[gridSize];
  this->boundary = new bool[gridSize];
  if (velocityX) {
    std::copy(velocityX, velocityX + gridSize, this->velocityX);
  } else {
    std::fill(this->velocityX, this->velocityX + gridSize, 0.0f);
  }
  if (velocityY) {
    std::copy(velocityY, velocityY + gridSize, this->velocityY);
  } else {
    std::fill(this->velocityY, this->velocityY + gridSize, 0.0f);
  }
  if (velocityZ) {
    std::copy(velocityZ, velocityZ + gridSize, this->velocityZ);
  } else {
    std::fill(this->velocityZ, this->velocityZ + gridSize, 0.0f);
  }
  std::copy(boundary, boundary + gridSize, this->boundary);

  /* Make sure the grid boundary is marked as such. */
  for (int i = 0; i < gridSizeX; i++) {
    for (int j = 0; j < gridSizeY; j++) {
      this->boundary[idx(i, j, 0)] = true;
      this->boundary[idx(i, j, gridSizeZ - 1)] = true;
    }
  }
  for (int i = 0; i < gridSizeX; i++) {
    for (int k = 0; k < gridSizeZ; k++) {
      this->boundary[idx(i, 0, k)] = true;
      this->boundary[idx(i, gridSizeY - 1, k)] = true;
    }
  }
  for (int j = 0; j < gridSizeY; j++) {
    for (int k = 0; k < gridSizeZ; k++) {
      this->boundary[idx(0, j, k)] = true;
      this->boundary[idx(gridSizeX - 1, j, k)] = true;
    }
  }
}

void StableINSSolver3D::updateVelocity(
    float *sourceX, float *sourceY, float *sourceZ,
    float viscosity, float vorticity, float dt,
    BoundaryCondition bc) {
  if (bc == BC_DIRICHLET) {
    copyBoundaryValues(velocityX, sourceX);
    copyBoundaryValues(velocityY, sourceY);
    copyBoundaryValues(velocityZ, sourceZ);

    copyBoundaryValues(tmpVelocityX, sourceX);
    copyBoundaryValues(tmpVelocityY, sourceY);
    copyBoundaryValues(tmpVelocityZ, sourceZ);
  }

  advect(tmpVelocityX, velocityX, dt, bc);
  advect(tmpVelocityY, velocityY, dt, bc);
  advect(tmpVelocityZ, velocityZ, dt, bc);

  addSource(tmpVelocityX, tmpVelocityX, sourceX, dt);
  addSource(tmpVelocityY, tmpVelocityY, sourceY, dt);
  addSource(tmpVelocityZ, tmpVelocityZ, sourceZ, dt);

  addVorticity(tmpVelocityX, tmpVelocityY, tmpVelocityZ, vorticity, dt, bc);

  diffuse(velocityX, tmpVelocityX, viscosity, dt, bc);
  diffuse(velocityY, tmpVelocityY, viscosity, dt, bc);
  diffuse(velocityZ, tmpVelocityZ, viscosity, dt, bc);

  float *pressure = tmpVelocityX;
  float *divergence = tmpVelocityY;
  projectVelocity(pressure, divergence, dt, bc);
}

void StableINSSolver3D::transportScalarField(
    float *field, float *source,
    float diffusion, float dissipation, float dt,
    BoundaryCondition bc) {
  if (bc == BC_DIRICHLET) {
    copyBoundaryValues(field, source);
  }

  addSource(field, field, source, dt);
  advect(source, field, dt, bc);
  diffuse(field, source, diffusion, dt, bc);
  dissipate(field, dissipation, dt);
}

void StableINSSolver3D::copyBoundaryValues(float *target, const float *source) {
  #pragma omp parallel for
  for (int i = 0; i < gridSize; i++) {
    if (boundary[i]) {
      target[i] = source[i];
    }
  }
}

void StableINSSolver3D::addSource(
    float *target,
    const float *field, const float *source,
    float dt) {
  #pragma omp parallel for
  for (int i = 0; i < gridSize; i++) {
    if (!boundary[i]) {
      target[i] = field[i] + dt * source[i];
    }
  }
}

void StableINSSolver3D::addVorticity(
    float *u, float *v, float *w,
    float vorticity, float dt,
    BoundaryCondition bc) {
  if (vorticity == 0.0f) {
    return;
  }

  if (!vorticityX) {
    vorticityX = new float[gridSize];
  }
  if (!vorticityY) {
    vorticityY = new float[gridSize];
  }
  if (!vorticityZ) {
    vorticityZ = new float[gridSize];
  }
  if (!tmpGrid) {
    tmpGrid = new float[gridSize];
  }

  float *curl = tmpGrid;

  #pragma omp parallel for
  for (int k = 1; k <= domainSizeZ; k++) {
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j, k)]) {
          float dwy =   value(w, i, j, k,   i  , j + 1,   k  , bc)
                      - value(w, i, j, k,   i  , j - 1,   k  , bc);
          float dvz =   value(v, i, j, k,   i  ,   j  , k + 1, bc)
                      - value(v, i, j, k,   i  ,   j  , k - 1, bc);
          float vx = dwy / (2.0f * hy) - dvz / (2.0f * hz);

          float duz =   value(u, i, j, k,   i  ,   j  , k + 1, bc)
                      - value(u, i, j, k,   i  ,   j  , k - 1, bc);
          float dwx =   value(w, i, j, k, i + 1,   j  ,   k  , bc)
                      - value(w, i, j, k, i - 1,   j  ,   k  , bc);
          float vy = duz / (2.0f * hz) - dwx / (2.0f * hx);

          float dvx =   value(v, i, j, k, i + 1,   j  ,   k  , bc)
                      - value(v, i, j, k, i - 1,   j  ,   k  , bc);
          float duy =   value(u, i, j, k,   i  , j + 1,   k  , bc)
                      - value(u, i, j, k,   i  , j - 1,   k  , bc);
          float vz = dvx / (2.0f * hx) - duy / (2.0f * hy);

          vorticityX[idx(i, j, k)] = vx;
          vorticityY[idx(i, j, k)] = vy;
          vorticityZ[idx(i, j, k)] = vz;

          curl[idx(i, j, k)] = std::sqrt(vx * vx + vy * vy + vz * vz);
        }
      }
    }
  }

  #pragma omp parallel for
  for (int k = 1; k <= domainSizeZ; k++) {
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j, k)]) {
          float nx =   value(curl, i, j, k, i + 1, j, k, BC_DIRICHLET_ZERO)
                     - value(curl, i, j, k, i - 1, j, k, BC_DIRICHLET_ZERO);
          nx /= 2.0f * hx;

          float ny =   value(curl, i, j, k, i, j + 1, k, BC_DIRICHLET_ZERO)
                     - value(curl, i, j, k, i, j - 1, k, BC_DIRICHLET_ZERO);
          ny /= 2.0f * hy;

          float nz =   value(curl, i, j, k, i, j, k + 1, BC_DIRICHLET_ZERO)
                     - value(curl, i, j, k, i, j, k - 1, BC_DIRICHLET_ZERO);
          nz /= 2.0f * hz;

          float nl = std::sqrt(nx * nx + ny * ny + nz * nz);
          if (nl > 0.00001f) {
            nx /= nl;
            ny /= nl;
            nz /= nl;

            float vx = vorticityX[idx(i, j, k)];
            float vy = vorticityY[idx(i, j, k)];
            float vz = vorticityZ[idx(i, j, k)];

            float fx = vorticity * hx * (ny * vz - vy * nz);
            float fy = vorticity * hy * (nz * vx - vz * nx);
            float fz = vorticity * hz * (nx * vy - vx * ny);

            u[idx(i, j, k)] += dt * fx;
            v[idx(i, j, k)] += dt * fy;
            w[idx(i, j, k)] += dt * fz;
          }
        }
      }
    }
  }
}

void StableINSSolver3D::diffuse(
    float *field, const float *field0,
    float factor, float dt,
    BoundaryCondition bc) {
  if (factor == 0.0f) {
    std::copy(field0, field0 + gridSize, field);
    return;
  }

  float ai = -dt * factor / (hx * hx);
  float aj = -dt * factor / (hy * hy);
  float ak = -dt * factor / (hz * hz);
  float a  = 1.0f - 2.0f * (ai + aj + ak);
  solveLinearSystem(field, field0, a, ai, aj, ak, bc);
}

void StableINSSolver3D::advect(
    float *field, const float *field0,
    float dt,
    BoundaryCondition bc) {
  #pragma omp parallel for
  for (int k = 1; k <= domainSizeZ; k++) {
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j, k)]) {
          float x = i - velocityX[idx(i, j, k)] * dt / hx;
          float y = j - velocityY[idx(i, j, k)] * dt / hy;
          float z = k - velocityZ[idx(i, j, k)] * dt / hz;

          float s1 = std::abs(x - (int) x);
          float s0 = 1.0f - s1;
          float t1 = std::abs(y - (int) y);
          float t0 = 1.0f - t1;
          float r1 = std::abs(z - (int) z);
          float r0 = 1.0f - r1;

          AdvectionTracker tracker(boundary, gridSizeX, gridSizeY, gridSizeZ);
          Bresenham::line3D(i, j, k, (int) x, (int) y, (int) z, tracker);
          int i0 = tracker.iEnd;
          int j0 = tracker.jEnd;
          int k0 = tracker.kEnd;

          float v1 = field0[idx(i0, j0, k0)];
          float v2 = value(field0, i0, j0, k0,   i0  ,   j0  , k0 + 1, bc);
          float v3 = value(field0, i0, j0, k0,   i0  , j0 + 1,   k0  , bc);
          float v4 = value(field0, i0, j0, k0,   i0  , j0 + 1, k0 + 1, bc);
          float v5 = value(field0, i0, j0, k0, i0 + 1,   j0  ,   k0  , bc);
          float v6 = value(field0, i0, j0, k0, i0 + 1,   j0  , k0 + 1, bc);
          float v7 = value(field0, i0, j0, k0, i0 + 1, j0 + 1,   k0  , bc);
          float v8 = value(field0, i0, j0, k0, i0 + 1, j0 + 1, k0 + 1, bc);

          float lerp1 = r0 * v1 + r1 * v2;
          float lerp2 = r0 * v3 + r1 * v4;
          float lerp3 = r0 * v5 + r1 * v6;
          float lerp4 = r0 * v7 + r1 * v8;

          float bilerp1 = t0 * lerp1 + t1 * lerp2;
          float bilerp2 = t0 * lerp3 + t1 * lerp4;

          float trilerp = s0 * bilerp1 + s1 * bilerp2;

          field[idx(i, j, k)] = trilerp;
        }
      }
    }
  }
}

void StableINSSolver3D::projectVelocity(
    float *pressure, float *divergence,
    float dt,
    BoundaryCondition bc) {
  #pragma omp parallel for
  for (int k = 1; k <= domainSizeZ; k++) {
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j, k)]) {
          float du =   value(velocityX, i, j, k, i + 1,   j  ,   k  , bc)
                     - value(velocityX, i, j, k, i - 1,   j  ,   k  , bc);
          float dv =   value(velocityY, i, j, k,   i  , j + 1,   k  , bc)
                     - value(velocityY, i, j, k,   i  , j - 1,   k  , bc);
          float dw =   value(velocityZ, i, j, k,   i  ,   j  , k + 1, bc)
                     - value(velocityZ, i, j, k,   i  ,   j  , k - 1, bc);

          divergence[idx(i, j, k)] = -(du / (2.0f * hx) + dv / (2.0f * hy) + dw / (2.0f * hz));

          pressure[idx(i, j, k)] = 0.0f;
        }
      }
    }
  }

  float ai = dt / (hx * hx);
  float aj = dt / (hy * hy);
  float ak = dt / (hz * hz);
  float a  = -2.0f * (ai + aj + ak);
  solveLinearSystem(pressure, divergence, -a, -ai, -aj, -ak, BC_NEUMANN_ZERO);

  #pragma omp parallel for
  for (int k = 1; k <= domainSizeZ; k++) {
    for (int j = 1; j <= domainSizeY; j++) {
      for (int i = 1; i <= domainSizeX; i++) {
        if (!boundary[idx(i, j, k)]) {
          float dp;

          dp =   value(pressure, i, j, k, i + 1, j, k, BC_NEUMANN_ZERO)
               - value(pressure, i, j, k, i - 1, j, k, BC_NEUMANN_ZERO);
          velocityX[idx(i, j, k)] -= dt * (dp / (2.0f * hx));

          dp =   value(pressure, i, j, k, i, j + 1, k, BC_NEUMANN_ZERO)
               - value(pressure, i, j, k, i, j - 1, k, BC_NEUMANN_ZERO);
          velocityY[idx(i, j, k)] -= dt * (dp / (2.0f * hy));

          dp =   value(pressure, i, j, k, i, j, k + 1, BC_NEUMANN_ZERO)
               - value(pressure, i, j, k, i, j, k - 1, BC_NEUMANN_ZERO);
          velocityZ[idx(i, j, k)] -= dt * (dp / (2.0f * hz));
        }
      }
    }
  }
}

void StableINSSolver3D::dissipate(float *field, float dissipation, float dt) {
  if (dissipation == 0.0f) {
    return;
  }

  #pragma omp parallel for
  for (int i = 0; i < gridSize; i++) {
    if (!boundary[i]) {
      field[i] /= 1.0f + dt * dissipation;
    }
  }
}

void StableINSSolver3D::solveLinearSystem(
    float *x, const float *x0,
    float a, float ai, float aj, float ak,
    BoundaryCondition bc) {
#if defined(JACOBI_SOLVER)
  if (!tmpGrid) {
    tmpGrid = new float[gridSize];
  }

  float *xNew = tmpGrid;

  if (bc == BC_DIRICHLET) {
    copyBoundaryValues(xNew, x);
  }

  for (int l = 0; l < 20; l++) {
    #pragma omp parallel for
    for (int k = 1; k <= domainSizeZ; k++) {
      for (int j = 1; j <= domainSizeY; j++) {
        for (int i = 1; i <= domainSizeX; i++) {
          if (!boundary[idx(i, j, k)]) {
            float sum  =   ai * value(x, i, j, k, i - 1,   j  ,   k  , bc)
                         + ai * value(x, i, j, k, i + 1,   j  ,   k  , bc)
                         + aj * value(x, i, j, k,   i  , j - 1,   k  , bc)
                         + aj * value(x, i, j, k,   i  , j + 1,   k  , bc)
                         + ak * value(x, i, j, k,   i  ,   j  , k - 1, bc)
                         + ak * value(x, i, j, k,   i  ,   j  , k + 1, bc);

            xNew[idx(i, j, k)] = (x0[idx(i, j, k)] - sum) / a;
          }
        }
      }
    }
    std::swap(x, xNew);
  }
#else
  for (int l = 0; l < 20; l++) {
    for (int k = 1; k <= domainSizeZ; k++) {
      for (int j = 1; j <= domainSizeY; j++) {
        for (int i = 1; i <= domainSizeX; i++) {
          if (!boundary[idx(i, j, k)]) {
            float sum  =   ai * value(x, i, j, k, i - 1,   j  ,   k  , bc)
                         + ai * value(x, i, j, k, i + 1,   j  ,   k  , bc)
                         + aj * value(x, i, j, k,   i  , j - 1,   k  , bc)
                         + aj * value(x, i, j, k,   i  , j + 1,   k  , bc)
                         + ak * value(x, i, j, k,   i  ,   j  , k - 1, bc)
                         + ak * value(x, i, j, k,   i  ,   j  , k + 1, bc);

            x[idx(i, j, k)] = (x0[idx(i, j, k)] - sum) / a;
          }
        }
      }
    }
  }
#endif
}

inline float StableINSSolver3D::value(
    const float *x,
    int i0, int j0, int k0,
    int i1, int j1, int k1,
    BoundaryCondition bc) {
  if (boundary[idx(i1, j1, k1)] && (bc != BC_DIRICHLET)) {
    if (bc == BC_NEUMANN_ZERO) {
      return x[idx(i0, j0, k0)];
    } else /* BC_DIRICHLET_ZERO */ {
      return 0.0f;
    }
  } else {
    return x[idx(i1, j1, k1)];
  }
}

inline int StableINSSolver3D::idx(int i, int j, int k) {
  return gridSizeX * (k * gridSizeY + j) + i;
}

void StableINSSolver3D::freeData() {
  delete[] velocityX;
  delete[] velocityY;
  delete[] velocityZ;
  delete[] boundary;
  delete[] vorticityX;
  delete[] vorticityY;
  delete[] vorticityZ;
  delete[] tmpGrid;
}

